package org.bjf.config;

import java.util.Map;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

import org.slf4j.MDC;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.sleuth.instrument.async.LazyTraceExecutor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.core.task.TaskExecutor;
import org.springframework.scheduling.annotation.AsyncConfigurer;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;

/**
 * @author binjinfeng
 */
@Configuration
@EnableAsync(proxyTargetClass = true)
@EnableScheduling
public class ThreadConfig implements AsyncConfigurer, SchedulingConfigurer {
  @Autowired
  private BeanFactory beanFactory;

  @Override
  public Executor getAsyncExecutor() {
    ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
    executor.setThreadNamePrefix("asyncExecutor-");
    executor.setCorePoolSize(8);
    executor.setMaxPoolSize(24);
    executor.setQueueCapacity(300);
    // 任务队列已满时的拒绝策略
    executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

    //线程池的优雅关闭
    executor.setWaitForTasksToCompleteOnShutdown(Boolean.TRUE);
    executor.setAwaitTerminationSeconds(60);
    //executor.setTaskDecorator(new ThreadLocalDecorator());
    executor.initialize();

    //return executor;
    // 使用sleuth的LazyTraceExecutor保证traceId在异步线程里传递
    return new LazyTraceExecutor(beanFactory,executor);
  }
  @Override
  public void configureTasks(ScheduledTaskRegistrar registrar) {
    ThreadPoolTaskScheduler taskScheduler=new ThreadPoolTaskScheduler();
    taskScheduler.setPoolSize(20);
    taskScheduler.setThreadNamePrefix("scheduler-");
    taskScheduler.setWaitForTasksToCompleteOnShutdown(Boolean.TRUE);
    taskScheduler.setAwaitTerminationSeconds(60);
    taskScheduler.initialize();

    registrar.setTaskScheduler(taskScheduler);
  }

  @Bean
  public TaskExecutor taskExecutor(){

    ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
    executor.setThreadNamePrefix("commExecutor-");
    executor.setCorePoolSize(5);
    executor.setMaxPoolSize(5);
    executor.setQueueCapacity(5);
    // 任务队列已满时的拒绝策略
    executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());

    //线程池的优雅关闭
    executor.setWaitForTasksToCompleteOnShutdown(Boolean.TRUE);
    executor.setAwaitTerminationSeconds(60);
    executor.initialize();

    return executor;
  }

  /**
   * 线程池里面传递threadLocal变量(示例)
   */
  private static class ThreadLocalDecorator implements TaskDecorator {

    @Override
    public Runnable decorate(Runnable runnable) {
      // 取出MDC变量
      Map<String, String> mdcContext = MDC.getCopyOfContextMap();
      return () -> {
        try {
          MDC.setContextMap(mdcContext);
          runnable.run();
        } finally {
          MDC.clear();
        }
      };
    }
  }
}
